// BWTandFM.cpp : Defines the entry point for the console application.
//

#include <fstream>
#include <iostream>
#include <string>
#include <vector>
#include <unordered_map>
#include <algorithm>
#include <cstring>
#include <memory>
#include <sys/resource.h>
#include <seqan/basic.h>
#include <seqan/file.h>
#include <seqan/sequence.h>
#include <seqan/translation.h>
#include <seqan/align.h>
#include <seqan/graph_align.h>
#include <seqan/graph_msa.h>
#include <seqan/graph_types.h>
#include <seqan/store.h>
#include <seqan/seq_io.h>
using namespace seqan;

const int alphabet = 100;

class IComponent {
public:
  virtual std::vector<int> operation(String<Dna>& str) = 0;
  virtual ~IComponent(){}
};

template <class T>
std::string convertToString(T t)
{
   std::stringstream ss;
   ss << t;
   std::string result = ss.str();
   return result;
}

class reverseComplementDecorator : public IComponent {
  

public:
  std::shared_ptr<IComponent> m_component;

  reverseComplementDecorator(IComponent* component) : m_component(component) {}

  virtual std::vector<int> operation(String<Dna>& str) {
    std::vector<int> firstVec = m_component->operation(str);
    reverseComplement(str);
    std::vector<int> revVec = m_component->operation(str);
    for (int i = 0; i < revVec.size(); ++i)
    {
      firstVec.push_back(revVec[i]);
    }
    return firstVec;
  }
};

struct FMIndex
{
  std::unordered_map<char, std::vector<int>> occ;
  std::unordered_map<char, int> c;
  FMIndex()
  {}

  FMIndex(std::vector<char> BWT)
  {
    std::vector<int> temp(alphabet);
    for (int i = 0; i < BWT.size(); ++i)
    {
      temp[BWT[i] + 1]++;
      occ[BWT[i]].resize(BWT.size());
    }
    for (int i = 1; i < temp.size(); ++i)
    {
      temp[i] += temp[i - 1];
      c[(char)i] = temp[i];

    }
    occ[BWT[0]][0] = 1;
    for (auto it : occ)
    {
      for (int i = 1; i < BWT.size(); ++i)
      {
        if (it.first == BWT[i])
          occ[it.first][i]++;
        occ[it.first][i] += occ[it.first][i - 1];
      }
    }
  }
};


class cBWT : public IComponent
{
private:
  FMIndex FM;
  std::vector<int> sArray;
  std::vector<char> BWT;
public:

  void constructFM()
  {
    FM = FMIndex(BWT);
  }

  std::vector<int> findOccures(std::string read)
  {
    int l = FM.c[read[read.length() - 1]];
    int r = FM.c[read[read.length() - 1] + 1] - 1;
    for (int i = read.length() - 2; i >= 0; --i)
    {
      if (l > r)
        break;
      l = FM.c[read[i]] + FM.occ[read[i]][l - 1];
      r = FM.c[read[i]] + FM.occ[read[i]][r] - 1;
    }
    std::vector<int> occs;
    for (int i = l; i <= r; ++i)
    {
      occs.push_back(sArray[i]);
    }
    return occs;
  }

  void constructSuffArray(std::string s)
  {
    const int maxlen = 1000000;
    int n = s.length();
    std::vector<int> p(n);

    int cnt[maxlen], c[maxlen];

    memset(cnt, 0, alphabet * sizeof(int));
    for (int i = 0; i<n; ++i)
      ++cnt[s[i]];
    for (int i = 1; i<alphabet; ++i)
      cnt[i] += cnt[i - 1];
    for (int i = 0; i<n; ++i)
      p[--cnt[s[i]]] = i;
    c[p[0]] = 0;
    int classes = 1;
    for (int i = 1; i<n; ++i) {
      if (s[p[i]] != s[p[i - 1]])  ++classes;
      c[p[i]] = classes - 1;
    }

    int pn[maxlen], cn[maxlen];
    for (int h = 0; (1 << h)<n; ++h) {
      for (int i = 0; i<n; ++i) {
        pn[i] = p[i] - (1 << h);
        if (pn[i] < 0)  pn[i] += n;
      }
      memset(cnt, 0, classes * sizeof(int));
      for (int i = 0; i<n; ++i)
        ++cnt[c[pn[i]]];
      for (int i = 1; i<classes; ++i)
        cnt[i] += cnt[i - 1];
      for (int i = n - 1; i >= 0; --i)
        p[--cnt[c[pn[i]]]] = pn[i];
      cn[p[0]] = 0;
      classes = 1;
      for (int i = 1; i<n; ++i) {
        int mid1 = (p[i] + (1 << h)) % n, mid2 = (p[i - 1] + (1 << h)) % n;
        if (c[p[i]] != c[p[i - 1]] || c[mid1] != c[mid2])
          ++classes;
        cn[p[i]] = classes - 1;
      }
      memcpy(c, cn, n * sizeof(int));
    }
    sArray = p;
  }

  void constructBWT(std::string& genome)
  {
    BWT.resize(genome.size());
    for (int i = 0; i < BWT.size(); ++i)
    {
      if (sArray[i] == 0)
      {
        BWT[i] = '$';
      }
      else
      {
        BWT[i] = genome[sArray[i] - 1];
      }
    }
  }

  cBWT(std::string& genome)
  {
    constructSuffArray(genome);
    constructBWT(genome);
    constructFM();
  }


  virtual std::vector<int> operation(String<Dna>& read)
  {
    return findOccures(convertToString(read));
  }

  virtual std::vector<int> operation(std::string& read)
  {
    return findOccures(read);
  }
};


void printOccs(String<Dna>& read, std::vector<int>& occur, std::ofstream& out)
{
  out << "Occurences of read " << read << " in genome are at positions : " << std::endl;
  for (int i = 0; i < occur.size(); ++i)
  {
    out << occur[i] << " ";
  }
  if (occur.size() == 0)
    out << "No entries";
  out << std::endl;
}

void usage()
{
   std::cerr << "Wrong params" << std::endl;
   std::cerr << "Usage: FMIndex reference.fasta reads.fasta" << std::endl;
}


int main(int argc, char* argv[])
{
  if (argc != 3)
  {
    usage();
    return 1;
  }

  const rlim_t kStackSize = 256 * 1024 * 1024;   // min stack size = 16 MB
  struct rlimit rl;
  int result;

  result = getrlimit(RLIMIT_STACK, &rl);
  if (result == 0)
  {
      if (rl.rlim_cur < kStackSize)
      {
          rl.rlim_cur = kStackSize;
          result = setrlimit(RLIMIT_STACK, &rl);
      }
  }


  std::ifstream genIn(argv[1]);
  String<Dna> seqgenome;
  read(genIn, seqgenome, Fasta());
  std::string genome = convertToString(seqgenome);
  genome.push_back('$');

  std::ifstream readIn(argv[2]);
  RecordReader<std::ifstream, SinglePass<> > reader(readIn);
  reverseComplementDecorator myBWT(new cBWT(genome));
  std::ofstream out("output.txt");
  std::cout << "Output will be written to output.txt" << std::endl;
  String<char> id;
  String<Dna> singleRead;
  while (!atEnd(reader))
  {
    if (readRecord(id, singleRead, reader, Fasta()) != 0)
    {
      std::cerr << "ERROR reading FASTA." << std::endl;
      return 1;
    }
    std::cout << singleRead << std::endl;
    std::vector<int> occur = myBWT.operation(singleRead);
    std::sort(occur.begin(), occur.end());
    printOccs(singleRead, occur, out);
  }
  return 0;
}

